# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
# coding=utf-8
# copyright (c) 2025 tencent inc. all rights reserved.
# xiaotaoliu@tencent.com, astrachang@tencent.com, nrwu@tencent.com, guanyouhe@tencent.com

from functools import partial

import torch

from megatron.core.transformer.moe.moe_utils import (
    switch_load_balancing_loss_func,
)

from mpatch.core.transformer.moe.moe_utils import topk_softmax_with_capacity


def aux_loss_load_balancing(self, logits: torch.Tensor):
    """Apply loss-based load balancing to the logits tensor.

    Args:
        logits (torch.Tensor): the logits tensor after gating, shape: [num_tokens, num_experts].

    Returns:
        probs (torch.Tensor): The probabilities of token to experts assignment.
        routing_map (torch.Tensor): The mask of token to experts assignment.
    """
    probs, routing_map, tokens_per_expert = topk_softmax_with_capacity(
        logits,
        self.topk,
        capacity_factor=self.config.moe_expert_capacity_factor,
        pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,
        drop_policy=self.config.moe_token_drop_policy,
        use_pre_softmax=self.config.moe_router_pre_softmax,
        num_groups=self.config.moe_router_num_groups,
        group_topk=self.config.moe_router_group_topk,
        scaling_factor=self.config.moe_router_topk_scaling_factor,
        deterministic_mode=self.config.deterministic_mode,
        score_function=self.score_function,
        expert_bias=self.expert_bias,
        # 传入 moe norm topk
        moe_norm_topk_prob=self.config.moe_norm_topk_prob,
        moe_norm_topk_prob_eps=self.config.moe_norm_topk_prob_eps,
    )

    if self.training:
        # Apply load balancing loss
        scores = torch.softmax(logits, dim=-1, dtype=torch.float32)
        aux_loss_func = partial(
            switch_load_balancing_loss_func,
            probs=scores,
            tokens_per_expert=tokens_per_expert,
            topk=self.topk,
        )
        probs = self.apply_load_balancing_loss(
            activation=probs, load_balancing_loss_func=aux_loss_func
        )
    return probs, routing_map
